import abc
from typing import TypeVar

from zkl_aiutils_datasets import Dataset, DelegatingDataset

from zkl_llmpt_iterator.tokenizer import TextTokenizer

AnySample = TypeVar('AnySample')


class LlmptDataset(DelegatingDataset[AnySample], abc.ABC):
    def __init__(self, dataset: Dataset[AnySample], tokenizer: TextTokenizer, total_tokens_n: int | None = None):
        super().__init__(dataset)
        self._tokenizer = tokenizer
        self._total_tokens_n = total_tokens_n

    @property
    def tokenizer(self) -> TextTokenizer:
        return self._tokenizer

    @property
    def total_tokens_n(self) -> int | None:
        return self._total_tokens_n
